import torch
from tempo import *
from tempo0 import *  # 含有setParameters等函数


def read_mean_variance_from_csvfile(csv_file):
    # 读取CSV文件
    df_result = pd.read_csv(csv_file)

    # 提取均值和方差
    mean = df_result['mean'].dropna().values
    variance = df_result['variance'].dropna().values

    # 将均值和方差转换为numpy.ndarray
    mean_array = np.array(mean)
    variance_array = np.array(variance)

    return mean_array, variance_array


def torch_cov(input_vec):
    """ 计算一个batch的协方差
    Args:
        input_vec: shape = (bn, seq, 2)
    """
    m  = input_vec - torch.mean(input_vec, dim=1, keepdim=True)
    mT = torch.transpose(m, 1, 2)
    cov_matrix = torch.matmul(mT, m) / (m.shape[1] - 1)
    return cov_matrix.reshape(input_vec.shape[0], -1)


def set_and_load_parameters_class_1(parameters):
    """设置和加载功能一的参数"""
    setParameters(parameters)
    return {"trajdata_path": parameters["trajdata_path"],
            "num_points": parameters["num_points"]}


def set_and_load_parameters_class_2(parameters):
    """设置和加载功能一的参数"""
    setParameters(parameters)
    return {"train_model_file": parameters["train_model_file"],
            'encoder_file': parameters['encoder_file'],
            "n_estimators": parameters["n_estimators"],
            "booster": parameters["booster"],
            'learning_rate': parameters['learning_rate'],
            "gamma": 0,
            'max_depth': parameters['max_depth']}


def set_and_load_parameters_class_3(parameters):
    """设置和加载功能一的参数"""
    setParameters(parameters)
    return {"pred_model_file": parameters["pred_model_file"],
            "encoder_file": parameters["encoder_file"]}

def set_and_load_parameters_imp_2(parameters_2):
    """设置和加载功能二的参数"""
    setParameters(parameters_2)
    return {"max_pair_targets": parameters_2['max_pair_targets'],
            'span_lower': parameters_2['span_lower'],
            'span_upper': parameters_2['span_upper'],
            'mask_factor': parameters_2['mask_factor'],
            'geometric_p': parameters_2['geometric_p'],
            'norm_file_path': parameters_2['norm_file_path']}


def set_and_load_parameters_imp_3(parameters_3):
    """设置和加载功能二的参数"""
    setParameters(parameters_3)
    return {"loc_size": parameters_3['loc_size'],
            'traj_len': parameters_3['traj_len'],
            'dim_in': parameters_3['dim_in'],
            'dim_out': parameters_3['dim_out'],
            'n_heads': parameters_3['n_heads'],
            'n_layers': parameters_3['n_layers'],
            'dropout': parameters_3['dropout'],
            'scale': parameters_3['scale'],
            'max_pair_targets': parameters_3['max_pair_targets'],
            'position_embedding_size': parameters_3['position_embedding_size'],
            'predict_model': parameters_3['predict_model'],
            'norm_file_path': parameters_3['norm_file_path'],
            'imp_save_path': parameters_3['imp_save_path']}



